import torch


def sinusoidal_positional_encoding(
    positions: torch.Tensor,
    periods: torch.Tensor,
) -> torch.Tensor:
    """ sinusoidal positional encoding
    :param positions: shape=[...]
    :param periods: shape=[periods_n]
    :return: encodings - shape=[..., periods_n*2]
    """
    positions = positions.unsqueeze(dim=-1)
    # [*pos_shape, 1]

    phase = positions / periods * 2 * torch.pi
    # [*pos_shape, chan_n //2]

    encodings_sin = torch.sin(phase)
    encodings_cos = torch.cos(phase)
    # [*pos_shape, chan_n //2]

    encodings = torch.cat([encodings_sin, encodings_cos], dim=-1)
    # [*pos_shape, chan_n]

    return encodings


def geometric_periods_sinusoidal_positional_encoding(
    positions: torch.Tensor, *,
    ratio: int | float | torch.Tensor = 2,
    period0: int | float | torch.Tensor = 2,
    periods_n: int,
) -> torch.Tensor:
    """ sinusoidal positional encoding whose periods is geometric sequence
    :param positions: shape=[...]
    :param ratio: float, the ratio of the geometric sequence, default=2.0
    :param period0: float, the head of the geometric sequence, default=2.0
    :param periods_n: int, the number of periods
    :return: encodings - shape=[..., periods_n*2]
    """
    ratio = torch.asarray(ratio, dtype=positions.dtype, device=positions.device)
    period0 = torch.asarray(period0, dtype=positions.dtype, device=positions.device)
    periods = torch.arange(0, periods_n, dtype=positions.dtype, device=positions.device)
    periods = torch.pow(ratio, periods) * period0
    return sinusoidal_positional_encoding(positions, periods)
